#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

def writeEcmpGroupRules(p4info_helper, switch, dst_ip, prefix_len, ecmp_base, ecmp_count):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_group",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip, prefix_len)
        },
        action_name="MyIngress.set_ecmp_select",
        action_params={
            "ecmp_base": ecmp_base,
            "ecmp_count": ecmp_count
        }
    )
    switch.WriteTableEntry(table_entry)
    print(f"Installed ECMP group rule: {dst_ip}/{prefix_len} -> base={ecmp_base}, count={ecmp_count}")

def writeEcmpNextHopRules(p4info_helper, switch, ecmp_index, nhop_dmac, nhop_ipv4, port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_nhop",
        match_fields={
            "meta.ecmp_select": ecmp_index
        },
        action_name="MyIngress.set_nhop",
        action_params={
            "nhop_dmac": nhop_dmac,
            "nhop_ipv4": nhop_ipv4,
            "port": port
        }
    )
    switch.WriteTableEntry(table_entry)
    print(f"Installed ECMP next hop: index={ecmp_index} -> MAC={nhop_dmac}, IP={nhop_ipv4}, port={port}")

def writeIpv4ForwardRules(p4info_helper, switch, dst_ip, prefix_len, dst_mac, port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip, prefix_len)
        },
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": dst_mac,
            "port": port
        }
    )
    switch.WriteTableEntry(table_entry)
    print(f"Installed IPv4 forward: {dst_ip}/{prefix_len} -> MAC={dst_mac}, port={port}")

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Initialize switch connections
        switches = {}
        switch_configs = [
            ('s1', '127.0.0.1:50051', 0),
            ('s2', '127.0.0.1:50052', 1),
            ('s3', '127.0.0.1:50053', 2),
            ('s4', '127.0.0.1:50054', 3),
            ('s5', '127.0.0.1:50055', 4),
            ('s6', '127.0.0.1:50056', 5)
        ]

        for name, address, device_id in switch_configs:
            switches[name] = p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name=name,
                address=address,
                device_id=device_id,
                proto_dump_file=f'logs/{name}-p4runtime-requests.txt'
            )
            switches[name].MasterArbitrationUpdate()
            switches[name].SetForwardingPipelineConfig(
                p4info=p4info_helper.p4info,
                bmv2_json_file_path=bmv2_file_path
            )
            print(f"Switch {name} connected and configured")

        # Install rules for s1 (handles traffic from h1 to VIP 10.0.0.1)
        print("\n=== Configuring Switch s1 ===")
        # ECMP group for VIP 10.0.0.1
        writeEcmpGroupRules(p4info_helper, switches['s1'], "10.0.0.1", 32, 0, 4)

        # ECMP next hops for VIP 10.0.0.1
        writeEcmpNextHopRules(p4info_helper, switches['s1'], 0, "00:00:00:00:01:02", "10.0.2.2", 2)  # to s2
        writeEcmpNextHopRules(p4info_helper, switches['s1'], 1, "00:00:00:00:01:03", "10.0.2.2", 3)  # to s3
        writeEcmpNextHopRules(p4info_helper, switches['s1'], 2, "00:00:00:00:01:04", "10.0.2.2", 4)  # to s4
        writeEcmpNextHopRules(p4info_helper, switches['s1'], 3, "00:00:00:00:01:05", "10.0.2.2", 5)  # to s5

        # Regular forwarding for h1
        writeIpv4ForwardRules(p4info_helper, switches['s1'], "10.0.1.1", 32, "08:00:00:00:01:01", 1)

        # Install rules for s6 (handles traffic from h2 to VIP 10.0.0.2)
        print("\n=== Configuring Switch s6 ===")
        # ECMP group for VIP 10.0.0.2
        writeEcmpGroupRules(p4info_helper, switches['s6'], "10.0.0.2", 32, 0, 4)

        # ECMP next hops for VIP 10.0.0.2
        writeEcmpNextHopRules(p4info_helper, switches['s6'], 0, "00:00:00:00:06:02", "10.0.1.1", 2)  # to s2
        writeEcmpNextHopRules(p4info_helper, switches['s6'], 1, "00:00:00:00:06:03", "10.0.1.1", 3)  # to s3
        writeEcmpNextHopRules(p4info_helper, switches['s6'], 2, "00:00:00:00:06:04", "10.0.1.1", 4)  # to s4
        writeEcmpNextHopRules(p4info_helper, switches['s6'], 3, "00:00:00:00:06:05", "10.0.1.1", 5)  # to s5

        # Regular forwarding for h2
        writeIpv4ForwardRules(p4info_helper, switches['s6'], "10.0.2.2", 32, "08:00:00:00:02:02", 1)

        # Install rules for intermediate switches (s2, s3, s4, s5)
        for switch_name in ['s2', 's3', 's4', 's5']:
            print(f"\n=== Configuring Switch {switch_name} ===")
            # Forward to h1 via s1
            writeIpv4ForwardRules(p4info_helper, switches[switch_name], "10.0.1.1", 32, "08:00:00:00:01:01", 1)
            # Forward to h2 via s6
            writeIpv4ForwardRules(p4info_helper, switches[switch_name], "10.0.2.2", 32, "08:00:00:00:02:02", 2)

        print("\n=== All rules installed successfully ===")
        print("Topology is ready for testing:")
        print("- h1 (10.0.1.1) -> VIP 10.0.0.1 -> ECMP via s2,s3,s4,s5 -> h2 (10.0.2.2)")
        print("- h2 (10.0.2.2) -> VIP 10.0.0.2 -> ECMP via s2,s3,s4,s5 -> h1 (10.0.1.1)")

        # Keep the controller running
        print("\nController running. Press Ctrl+C to exit.")
        while True:
            sleep(1)

    except KeyboardInterrupt:
        print("\nShutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    finally:
        ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Load Balance Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)

    main(args.p4info, args.bmv2_json)
